#include "getAmmo_cv.h"

using namespace cv;
using namespace std;

FindAmmo::FindAmmo()
{
	kpts1.reserve(56);
	matcher.create(NORM_HAMMING);
	FileStorage fs("kp.xml", FileStorage::READ);
	for (size_t i = 0; i < 56; i++)
	{
		fs["kpts" + to_string(i)] >> kpts1[i];
	}
	fs["desc"] >> desc1;
	fs.release();
	akaze->setNOctaves(3);
	akaze->setNOctaveLayers(4);
	akaze->setDescriptorType(AKAZE::DESCRIPTOR_MLDB_UPRIGHT);
	matched1.reserve(56);
	matched2.reserve(56);
	nn_matches.reserve(56);
	inlier_matches.reserve(56);
}

vector<Point2f> FindAmmo::Points(vector<KeyPoint> keypoints, Point2f bias = Point(0, 0))
{
	vector<Point2f> res;
	for (unsigned i = 0; i < keypoints.size(); i++)
	{
		res.push_back(keypoints[i].pt + bias);
	}
	return res;
}

void FindAmmo::ReleaseAll()
{
	imshow("result", img2);
	waitKey(1);
	matched1.clear();
	matched2.clear();
	inliers1.clear();
	inliers2.clear();
	nn_matches.clear();
	inlier_matches.clear();
}

void FindAmmo::DetectNothing()
{
	cout << "*******************************" << endl;
	cout << "# Keypoints 1:                        \t" << kpts1.size() << endl;
	cout << "# Keypoints 2:                        \t" << kpts2.size() << endl;
	cout << "# Matches:                            \t" << matched1.size() << endl;
	cout << endl;
	ReleaseAll();
}

int FindAmmo::getAmmo(Mat frame)
{
	img2 = frame;

	double t = (double)getTickCount();
	Mat img2_ROI = img2(ROI);
	cvtColor(img2_ROI, img2_ROI, CV_BGR2GRAY);
	akaze->detectAndCompute(img2_ROI, noArray(), kpts2, desc2);
	if (kpts2.empty())
	{
		DetectNothing();
		return 0;
	}
	rectangle(img2, ROI, Scalar(0, 255, 0));
	matcher.knnMatch(desc1, desc2, nn_matches, 2);
	for (size_t i = 0; i < nn_matches.size(); i++)
	{
		DMatch first = nn_matches[i][0];
		float dist1 = nn_matches[i][0].distance;
		float dist2 = nn_matches[i][1].distance;
		if (dist1 < 0.8 * dist2)
		{
			matched1.push_back(kpts1[first.queryIdx]);
			matched2.push_back(kpts2[first.trainIdx]);
		}
	}
	Mat inlier_mask, homography;
	if (matched1.size() >= 10)
	{
		homography = findHomography(Points(matched1), Points(matched2, ROI.tl()),
									RANSAC, 3, inlier_mask, 2000, 0.9);
	}
	else
	{
		DetectNothing();
		return 0;
	}

	int inlier_cnt = 0;
	for (size_t i = 0; i < inlier_mask.rows; i++)
	{
		if (inlier_mask.at<uchar>(i))
			inlier_cnt++;
	}

	if (inlier_cnt < 10 || homography.empty())
	{
		DetectNothing();
		return 0;
	}

	/*vector<Point2f> bb = { Point2f(640 / 2 + 50,480 / 2 - 50),Point2f(640 / 2 - 50,480 / 2 - 50),
			Point2f(640 / 2 - 50,480 / 2 + 50),Point2f(640 / 2 + 50,480 / 2 + 50) };
		vector<Point2f> new_bb;
		perspectiveTransform(bb, new_bb, homography);
		line(img2, new_bb[0], new_bb[1], Scalar(255, 255, 0));
		line(img2, new_bb[1], new_bb[2], Scalar(255, 0, 255));
		line(img2, new_bb[2], new_bb[3], Scalar(0, 255, 255));
		line(img2, new_bb[3], new_bb[0], Scalar(255, 127, 127));*/

	double *homo0 = homography.ptr<double>(0);
	double *homo1 = homography.ptr<double>(1);
	double *homo2 = homography.ptr<double>(2);
	double scale = homo2[0] + homo2[1] + homo2[2];
	double trans_X = (homo0[0] + homo0[1]) / scale;
	double trans_Y = (homo1[0] + homo1[1]) / scale;
	//putText(img2, to_string(trans_X), Point(100, 100), 2, 2, Scalar(255, 0, 0), 2);
	//putText(img2, to_string(trans_Y), Point(100, 300), 2, 2, Scalar(255, 0, 0), 2);
	if (trans_X < 0.1 || trans_X > 3.00 || trans_Y < 0.1 || trans_Y > 3.00)
	{
		DetectNothing();
		return 0;
	}

	double new_X = (homo0[0] * 320 + homo0[1] * 240 + homo0[2]) / (homo2[0] * 320 + homo2[1] * 240 + homo2[2]);
	cout << endl
		 << "x:" << new_X-320 << endl;

	t = ((double)getTickCount() - t) / getTickFrequency();
	cout << "Times:" << int(t*1000) << endl;
	line(img2, Point(new_X, 0), Point(new_X, 479), Scalar(255), 2);

	if (new_X < 90)
		new_X = 90;
	if (new_X > 439)
		new_X = 439;
	ROI = Rect(Point(0, 0), Point(640, 480));

	ReleaseAll();

#ifdef _DEBUG_
	for (unsigned i = 0; i < matched1.size(); i++)
	{
		if (inlier_mask.at<uchar>(i))
		{
			int new_i = static_cast<int>(inliers1.size());
			inliers1.push_back(matched1[i]);
			inliers2.push_back(matched2[i]);
			inlier_matches.push_back(DMatch(new_i, new_i, 0));
		}
	}

	Mat res;
	Mat img1 = imread("img.jpg", IMREAD_GRAYSCALE);
	drawMatches(img1, inliers1, img2, inliers2, inlier_matches, res);
	double inlier_ratio = inliers1.size() * 1.0 / matched1.size();
	cout << "*******************************" << endl;
	cout << "# Keypoints 1:                        \t" << kpts1.size() << endl;
	cout << "# Keypoints 2:                        \t" << kpts2.size() << endl;
	cout << "# Matches:                            \t" << matched1.size() << endl;
	cout << "# Inliers:                            \t" << inliers1.size() << endl;
	cout << "# Inliers Ratio:                      \t" << inlier_ratio << endl;
	cout << endl;
	imshow("result", res);
	if (waitKey(1) == 'q')
		break;
#endif

	return 0;
}
